{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Exploring the Impact of Evaluation Order on Edit Distance Algorithms\n",
                "\n",
                "Removing data-dependencies in the Wagner-Fisher, Needleman-Wunsch, Smith-Waterman, and Gotoh Dynamic Programming algorithms to explain the hardware-accelerated variants in StringZilla."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Levenshtein Distance\n",
                "\n",
                "Levenshtein edit distance is one of the most broadly studied string similarity metrics.\n",
                "It is defined as the minimum number of single-character insertions, deletions, and substitutions required to change one string into another.\n",
                "The Levenshtein distance between two strings is calculated using dynamic programming algorithms, such as the Wagner-Fisher algorithm, and its variations for Bioinformatics: \n",
                "\n",
                "- Needleman-Wunsch for global alignment with substitution matrices, \n",
                "- Smith-Waterman for local alignment with substitution matrices, \n",
                "- Gotoh for different penalties for gap opening and extensions.\n",
                "\n",
                "Given the shared nature of these algorithms, the same tricks can be applied to all of them to improve their performance."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Warner-Fisher Algorithm\n",
                "\n",
                "Wagner-Fisher algorithm, in its most naive form, has a time and space complexity of $O(NM)$, where $N$ and $M$ are the lengths of the two strings being compared.\n",
                "A rectangular matrix of size $(N+1) \\times (M+1)$ is created to store the edit distances between all prefixes of the two strings.\n",
                "The first row and column are, naturally, initialized with ${0, 1, 2, ..., N}$ and ${0, 1, 2, ..., M}$ respectively."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from typing import Tuple\n",
                "import numpy as np # NumPy for matrices\n",
                "\n",
                "def wagner_fisher(s1: str, s2: str) -> Tuple[int, np.ndarray]:\n",
                "    # Create a matrix of size (len(s1)+1) x (len(s2)+1)\n",
                "    matrix = np.zeros((len(s1) + 1, len(s2) + 1), dtype=int)\n",
                "\n",
                "    # Initialize the first column and first row of the matrix\n",
                "    for i in range(len(s1) + 1):\n",
                "        matrix[i, 0] = i\n",
                "    for j in range(len(s2) + 1):\n",
                "        matrix[0, j] = j\n",
                "\n",
                "    # Compute Levenshtein distance\n",
                "    for i in range(1, len(s1) + 1):\n",
                "        for j in range(1, len(s2) + 1):\n",
                "            substitution_cost = s1[i - 1] != s2[j - 1]\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      #? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      #? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  #? Substitution cost\n",
                "            )\n",
                "\n",
                "    # The distance will be placed in the bottom right corner of the matrix\n",
                "    return matrix[len(s1), len(s2)], matrix"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "s1 = \"kiten\"\n",
                "s2 = \"katerinas\"\n",
                "distance_wf, matrix_wf = wagner_fisher(s1, s2)\n",
                "s1, s2, f\"{distance_wf = }\", matrix_wf"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "This algorithm is almost never recommended for practical use, as it has a quadratic space complexity.\n",
                "It's trivial to see that the space complexity can be reduced to $O(min(N, M))$ by only storing the last two rows of the matrix, but we want to keep the entire matrix as a reference to allow debugging and visualization."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "To feel safer, while designing our alternative traversal algorithm, let's define an extraction function, that will get the values of a certain skewed diagonal."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def get_skewed_diagonal(matrix: np.ndarray, index: int):\n",
                "    flipped_matrix = np.fliplr(matrix)\n",
                "    return np.flip(np.diag(flipped_matrix, k= matrix.shape[1] - index - 1))\n",
                "\n",
                "# Let's test this function right away.\n",
                "matrix = np.array([\n",
                "    [1, 2, 3],\n",
                "    [4, 5, 6],\n",
                "    [7, 8, 9]])\n",
                "assert np.all(get_skewed_diagonal(matrix, 2) == [7, 5, 3])\n",
                "assert np.all(get_skewed_diagonal(matrix, 1) == [4, 2])\n",
                "assert np.all(get_skewed_diagonal(matrix, 4) == [9])\n",
                "\n",
                "# Let's test this function right away.\n",
                "matrix = np.array([\n",
                "    [1, 2, 3],\n",
                "    [4, 5, 6]])\n",
                "assert np.all(get_skewed_diagonal(matrix, 0) == [1])\n",
                "assert np.all(get_skewed_diagonal(matrix, 1) == [4, 2])\n",
                "assert np.all(get_skewed_diagonal(matrix, 2) == [5, 3])\n",
                "assert np.all(get_skewed_diagonal(matrix, 3) == [6])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "get_skewed_diagonal(matrix_wf, 10)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Diagonal Evaluation Order\n",
                "\n",
                "Accelerating this exact algorithm with SIMD instructions isn't trivial, is the `matrix[i, j]` value has a dependency on the `matrix[i, j - 1]` value.\n",
                "So we can't brute-force accelerate the inner loop.\n",
                "Instead, we can show that we can evaluate the matrix in a different order, and still get the same result.\n",
                "\n",
                "![Skewed Diagonals Evaluation Order](https://mathworld.wolfram.com/images/eps-svg/SkewDiagonal_1000.svg)\n",
                "\n",
                "But before complicating things too much, let's start with a simple case - when both strings have identical lengths and the DP matrix has a square shape."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from typing import Optional\n",
                "\n",
                "def square_skewed_diagonals(\n",
                "    s1: str, s2: str, \n",
                "    verbose: bool = False, \n",
                "    baseline: Optional[np.ndarray] = None) -> Tuple[int, np.ndarray]:\n",
                "\n",
                "    assert len(s1) == len(s2), \"First define an algo for square matrices!\"\n",
                "    # Create a matrix of size (len(s1)+1) x (len(s2)+1)\n",
                "    matrix = np.zeros((len(s1) + 1, len(s2) + 1), dtype=int)\n",
                "    matrix[:, :] = 99\n",
                "\n",
                "    # Initialize the first column and first row of the matrix\n",
                "    for i in range(len(s1) + 1):\n",
                "        matrix[i, 0] = i\n",
                "    for j in range(len(s2) + 1):\n",
                "        matrix[0, j] = j\n",
                "\n",
                "    # Number of rows and columns in the square matrix.\n",
                "    n = len(s1) + 1\n",
                "    \n",
                "    # Number of diagonals and skewed diagonals in the square matrix of size (n x n).\n",
                "    diagonals_count = 2 * n - 1\n",
                "    \n",
                "    # Populate the matrix in 2 separate loops: for the top left triangle and for the bottom right triangle.\n",
                "    for skew_diagonal_index in range(2, n):\n",
                "        skew_diagonal_length = skew_diagonal_index + 1\n",
                "        for offset_within_diagonal in range(1, skew_diagonal_length - 1):\n",
                "            # If we haven't passed the main skew diagonal yet, \n",
                "            # then we have to skip the first and the last operation,\n",
                "            # as those are already pre-populated and form the first column \n",
                "            # and the first row of the Levenshtein matrix respectively.\n",
                "            i = skew_diagonal_index - offset_within_diagonal\n",
                "            j = offset_within_diagonal\n",
                "            if verbose:\n",
                "                print(f\"top left triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            substitution_cost = s1[i - 1] != s2[j - 1]\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      #? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      #? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  #? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "            \n",
                "    # Now the bottom right triangle of the matrix.\n",
                "    for skew_diagonal_index in range(n, diagonals_count):\n",
                "        skew_diagonal_length = 2 * n - skew_diagonal_index - 1\n",
                "        for offset_within_diagonal in range(skew_diagonal_length):\n",
                "            i = n - offset_within_diagonal - 1\n",
                "            j = skew_diagonal_index - n + offset_within_diagonal + 1\n",
                "            if verbose:\n",
                "                print(f\"bottom right triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            substitution_cost = s1[i - 1] != s2[j - 1]\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      #? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      #? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  #? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "\n",
                "    # Similarly, the distance will be placed in the bottom right corner of the matrix\n",
                "    return matrix[len(s1), len(s2)], matrix"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Let's generate some random strings and make sure we produce the right result."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import random\n",
                "for _ in range(10):\n",
                "    s1 = ''.join(random.choices(\"abc\", k=50))\n",
                "    s2 = ''.join(random.choices(\"abc\", k=50))\n",
                "    distance_wf, matrix_wf = wagner_fisher(s1, s2)\n",
                "    distance_sd, matrix_sd = square_skewed_diagonals(s1, s2, baseline=matrix_wf)\n",
                "    assert distance_wf == distance_sd, f\"{distance_wf = } != {distance_sd = }\"\n",
                "    assert np.all(matrix_wf == matrix_sd), f\"{matrix_wf = }\\n{matrix_sd = }\""
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Vectorizing the Skewed Diagonals Algorithm\n",
                "\n",
                "Going further, we can avoid storing the whole matrix, and only store three diagonals at a time.\n",
                "The longer will never exceed `n` in length.\n",
                "The others are always at most `n-1`.\n",
                "Let's try vectorizing different parts of our algorithm, validating it against the output of the naive algorithm for 2 strings: `\"BCDE\"` and `\"FKPU\"`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "s1 = \"BCDE\"\n",
                "s2 = \"FKPU\"\n",
                "distance_wf, matrix_wf = wagner_fisher(s1, s2)\n",
                "s1, s2, f\"{distance_wf = }\", matrix_wf"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Replacing the letters with numbers and annotating with a header row and column for `\"BCDE\"` and `\"FKPU\"`:\n",
                "\n",
                "|       |     | **B** | **C** | **D** | **E** |\n",
                "| ----- | --- | ----- | ----- | ----- | ----- |\n",
                "|       | a   | b     | c     | d     | e     |\n",
                "| **F** | f   | g     | h     | i     | j     |\n",
                "| **K** | k   | l     | m     | n     | o     |\n",
                "| **P** | p   | q     | r     | s     | t     |\n",
                "| **U** | u   | v     | w     | x     | y     |\n",
                "\n",
                "At any point we will be working with 3 diagonals:\n",
                "\n",
                "- `previous` set to `[a]` at start\n",
                "- `current` set to `[f, b]` at start\n",
                "- `following` set to `[k, g, c]` at start"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "assert len(s1) == len(s2), \"First define an algo for square matrices!\"\n",
                "# Number of rows and columns in the square matrix.\n",
                "n = len(s1) + 1\n",
                "\n",
                "following = np.zeros(n, dtype=np.uint) # let's assume we are computing the main skew diagonal: [u, q, m, i, e]\n",
                "current = np.zeros(n, dtype=np.uint) # will contain: [p, l, h, e]\n",
                "previous = np.zeros(n, dtype=np.uint) # will contain: [k, g, c]\n",
                "\n",
                "# Initialize the first two diagonals.\n",
                "# The `previous` would contain the values [a].\n",
                "# The `current` would contain the values [f, b]. \n",
                "previous[0] = 0\n",
                "current[0:2] = 1\n",
                "previous, current, following"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now we can rewrite the first nested loop for the upper-left triangle of the matrix in NumPy primitives, using it's `np.minimum` function to calculate the minimum of three values."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# To evaluate every subsequent entry:\n",
                "next_diagonal_index = 2\n",
                "while next_diagonal_index < n:\n",
                "    next_skew_diagonal_length = next_diagonal_index + 1\n",
                "\n",
                "    old_substitution_costs = previous[:next_skew_diagonal_length - 2]\n",
                "    added_substitution_costs = [s1[next_diagonal_index - i - 2] != s2[i] for i in range(next_skew_diagonal_length - 2)]\n",
                "    substitution_costs = old_substitution_costs + added_substitution_costs\n",
                "\n",
                "    following[1:next_skew_diagonal_length - 1] = np.minimum(current[1:next_skew_diagonal_length - 1] + 1, current[:next_skew_diagonal_length - 2] + 1) # Insertions or deletions\n",
                "    following[1:next_skew_diagonal_length - 1] = np.minimum(following[1:next_skew_diagonal_length - 1], substitution_costs) # Substitutions\n",
                "    following[0] = next_diagonal_index\n",
                "    following[next_skew_diagonal_length - 1] = next_diagonal_index\n",
                "    assert np.all(following[:next_skew_diagonal_length] == get_skewed_diagonal(matrix_wf, next_diagonal_index))\n",
                "    \n",
                "    previous[:] = current[:]\n",
                "    current[:] = following[:]\n",
                "    next_diagonal_index += 1\n",
                "\n",
                "previous, current, following # Log the state"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "By now we've scanned through the upper-left triangle of the matrix, where each subsequent iteration results in a larger diagonal.\n",
                "From now onwards, we will be shrinking.\n",
                "Instead of adding value equal to the skewed diagonal index on either side, we will be cropping those values out."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "while next_diagonal_index < 2 * n - 1:\n",
                "    next_skew_diagonal_length = 2 * n - 1 - next_diagonal_index\n",
                "    old_substitution_costs = previous[:next_skew_diagonal_length]\n",
                "    added_substitution_costs = [s1[len(s1) - i - 1] != s2[next_diagonal_index - n + i] for i in range(next_skew_diagonal_length)]\n",
                "    substitution_costs = old_substitution_costs + added_substitution_costs\n",
                "    \n",
                "    following[:next_skew_diagonal_length] = np.minimum(current[:next_skew_diagonal_length] + 1, current[1 : next_skew_diagonal_length + 1] + 1) # Insertions or deletions\n",
                "    following[:next_skew_diagonal_length] = np.minimum(following[:next_skew_diagonal_length], substitution_costs) # Substitutions\n",
                "    assert np.all(following[:next_skew_diagonal_length] == get_skewed_diagonal(matrix_wf, next_diagonal_index)), f\"\\n{following[:next_skew_diagonal_length]} not equal to \\n{get_skewed_diagonal(baseline, next_diagonal_index)}\"\n",
                "    \n",
                "    previous[:next_skew_diagonal_length] = current[1:next_skew_diagonal_length + 1]\n",
                "    current[:next_skew_diagonal_length] = following[:next_skew_diagonal_length]\n",
                "    next_diagonal_index += 1\n",
                "\n",
                "previous, current, following # Log the state"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "assert distance_wf == following[0], f\"{distance_wf = } != {following[0] = }\""
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Generalizing to Non-Square Matrices\n",
                "\n",
                "Let's imaging 2 inputs of length 3 and 5: `\"KPU\"` and `\"BCDEF\"`:\n",
                "\n",
                "|       |     | **B** | **C** | **D** | **E** | **F** |\n",
                "| ----- | --- | ----- | ----- | ----- | ----- | ----- |\n",
                "|       | a   | b     | c     | d     | e     | f     |\n",
                "| **K** | g   | h     | i     | j     | k     | l     |\n",
                "| **P** | m   | n     | o     | p     | q     | r     |\n",
                "| **U** | s   | t     | u     | v     | w     | x     |\n",
                "\n",
                "At any point we will be working with 3 diagonals:\n",
                "\n",
                "- `previous` set to `[a]` at start\n",
                "- `current` set to `[g, b]` at start\n",
                "- `next` set to `[m, h, c]` at start\n",
                "\n",
                "Once we proceed to for X cycles:\n",
                "\n",
                "- `previous` set to `[s, n, i, d]`\n",
                "- `current` set to `[t, o, j, e]`\n",
                "- `next` set to `[u, p, k, f]`\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from typing import Optional\n",
                "\n",
                "def skewed_diagonals(\n",
                "    s1: str, s2: str, \n",
                "    verbose: bool = False, \n",
                "    baseline: Optional[np.ndarray] = None) -> Tuple[int, np.ndarray]:\n",
                "    \n",
                "    shorter, longer = (s1, s2) if len(s1) <= len(s2) else (s2, s1)    \n",
                "    baseline = baseline if len(s1) <= len(s2) else baseline.T\n",
                "    shorter_dim = len(shorter) + 1\n",
                "    longer_dim = len(longer) + 1\n",
                "    if verbose:\n",
                "        print(f\"{shorter=}, {longer=}, {shorter_dim=}, {longer_dim=}\")\n",
                "    \n",
                "    # Create a matrix of size (shorter_dim) x (longer_dim)\n",
                "    matrix = np.zeros((shorter_dim, longer_dim), dtype=int)\n",
                "    matrix[:, :] = longer_dim + 1 # or +inf \n",
                "\n",
                "    # Initialize the first column and first row of the matrix\n",
                "    for i in range(shorter_dim):\n",
                "        matrix[i, 0] = i\n",
                "    for j in range(longer_dim):\n",
                "        matrix[0, j] = j\n",
                "\n",
                "    # Let's say we are dealing with 3 and 5 letter words.\n",
                "    # The matrix will have size 4 x 6, parameterized as (shorter_dim x longer_dim).\n",
                "    # It will have:\n",
                "    # - 4 diagonals of increasing length, at positions: 0, 1, 2, 3.\n",
                "    # - 2 diagonals of fixed length, at positions: 4, 5.\n",
                "    # - 3 diagonals of decreasing length, at positions: 6, 7, 8.\n",
                "    diagonals_count = shorter_dim + longer_dim - 1\n",
                "\n",
                "    # Same as with square matrices, the 0th diagonal contains - just one element - zero - skipping it.\n",
                "    # Same as with square matrices, the 1st diagonal contains the values 1 and 1 - skipping it.\n",
                "    # Now let's handle the rest of the upper-left triangle.\n",
                "    for skew_diagonal_index in range(2, shorter_dim):\n",
                "        skew_diagonal_length = (skew_diagonal_index + 1)\n",
                "        for offset_within_diagonal in range(1, skew_diagonal_length - 1): # ! Skip the first column & row\n",
                "            # If we haven't passed the main skew diagonal yet, \n",
                "            # then we have to skip the first and the last operation,\n",
                "            # as those are already pre-populated and form the first column \n",
                "            # and the first row of the Levenshtein matrix respectively.\n",
                "            i = skew_diagonal_index - offset_within_diagonal\n",
                "            j = offset_within_diagonal\n",
                "            if verbose:\n",
                "                print(f\"top left triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            shorter_char = shorter[i - 1]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "                            \n",
                "    # Now let's handle the anti-diagonal band of the matrix, between the top and bottom-right triangles.        \n",
                "    for skew_diagonal_index in range(shorter_dim, longer_dim):\n",
                "        skew_diagonal_length = shorter_dim\n",
                "        for offset_within_diagonal in range(skew_diagonal_length - 1): # ! Skip the first row\n",
                "            i = shorter_dim - offset_within_diagonal - 1\n",
                "            j = skew_diagonal_index - shorter_dim + offset_within_diagonal + 1\n",
                "            if verbose:\n",
                "                print(f\"anti-band: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            shorter_char = shorter[i - 1]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "    \n",
                "    # Now let's handle the bottom right triangle.\n",
                "    for skew_diagonal_index in range(longer_dim, diagonals_count):\n",
                "        skew_diagonal_length = diagonals_count - skew_diagonal_index\n",
                "        for offset_within_diagonal in range(skew_diagonal_length):\n",
                "            i = shorter_dim - offset_within_diagonal - 1\n",
                "            j = skew_diagonal_index - shorter_dim + offset_within_diagonal + 1\n",
                "            if verbose:\n",
                "                print(f\"bottom right triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            assert (i - 1) >= 0 and (i - 1) < len(shorter), f\"{i = }\"\n",
                "            assert (j - 1) >= 0 and (j - 1) < len(longer), f\"{j = }\"\n",
                "            shorter_char = shorter[i - 1]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "\n",
                "    # Return the Levenshtein distance\n",
                "    distance = matrix[len(shorter), len(longer)]\n",
                "    if len(s1) > len(s2):\n",
                "        matrix = matrix.T\n",
                "    return distance, matrix"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import random\n",
                "for _ in range(100):\n",
                "    len1 = random.randint(1, 50)\n",
                "    len2 = random.randint(1, 50)\n",
                "    s1 = ''.join(random.choices(\"abc\", k=len1))\n",
                "    s2 = ''.join(random.choices(\"abc\", k=len2))\n",
                "    distance_wf, matrix_wf = wagner_fisher(s1, s2)\n",
                "    distance_sd, matrix_sd = skewed_diagonals(s1, s2, baseline=matrix_wf, verbose=False)\n",
                "    assert distance_wf == distance_sd, f\"{distance_wf = } != {distance_sd = }\"\n",
                "    assert np.all(matrix_wf == matrix_sd), f\"{matrix_wf = }\\n{matrix_sd = }\""
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "s1 = \"listeners\"\n",
                "s2 = \"silents\"\n",
                "distance_wf, matrix_wf = wagner_fisher(s1, s2)\n",
                "distance_sd, matrix_sd = skewed_diagonals(s1, s2, baseline=matrix_wf)\n",
                "s1, s2, f\"{distance_sd = }\", matrix_sd"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Reversing the Input\n",
                "\n",
                "One of the issues with vectorizing this algorithm is the traversal order of the shorter string.\n",
                "It's different from the longer string and different from the natural traversal order of the loop.\n",
                "To make the indexing simpler, we can pre-reverse the shorter string."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from typing import Optional\n",
                "\n",
                "def skewed_diagonals_reversed(\n",
                "    s1: str, s2: str, \n",
                "    verbose: bool = False, \n",
                "    baseline: Optional[np.ndarray] = None) -> Tuple[int, np.ndarray]:\n",
                "    \n",
                "    shorter, longer = (s1, s2) if len(s1) <= len(s2) else (s2, s1)    \n",
                "    baseline = baseline if len(s1) <= len(s2) else baseline.T\n",
                "    shorter_dim = len(shorter) + 1\n",
                "    longer_dim = len(longer) + 1\n",
                "    if verbose:\n",
                "        print(f\"{shorter=}, {longer=}, {shorter_dim=}, {longer_dim=}\")\n",
                "    \n",
                "    # Create a matrix of size (shorter_dim) x (longer_dim)\n",
                "    matrix = np.zeros((shorter_dim, longer_dim), dtype=int)\n",
                "    matrix[:, :] = longer_dim + 1 # or +inf \n",
                "\n",
                "    # Initialize the first column and first row of the matrix\n",
                "    for i in range(shorter_dim):\n",
                "        matrix[i, 0] = i\n",
                "    for j in range(longer_dim):\n",
                "        matrix[0, j] = j\n",
                "\n",
                "    # Let's say we are dealing with 3 and 5 letter words.\n",
                "    # The matrix will have size 4 x 6, parameterized as (shorter_dim x longer_dim).\n",
                "    # It will have:\n",
                "    # - 4 diagonals of increasing length, at positions: 0, 1, 2, 3.\n",
                "    # - 2 diagonals of fixed length, at positions: 4, 5.\n",
                "    # - 3 diagonals of decreasing length, at positions: 6, 7, 8.\n",
                "    diagonals_count = shorter_dim + longer_dim - 1\n",
                "    shorter_reversed = \"\".join(reversed(shorter))\n",
                "\n",
                "    # In reality, we need to keep only 3 diagonals to produce the same score in the end.\n",
                "    previous_distances = np.zeros(shorter_dim, dtype=np.uint)\n",
                "    current_distances = np.zeros(shorter_dim, dtype=np.uint)\n",
                "    next_distances = np.zeros(shorter_dim, dtype=np.uint)\n",
                "    temporary_distances = np.zeros(shorter_dim, dtype=np.uint)\n",
                "    previous_distances[0] = 0\n",
                "    current_distances[0] = current_distances[1] = 1\n",
                "\n",
                "    # Same as with square matrices, the 0th diagonal contains - just one element - zero - skipping it.\n",
                "    # Same as with square matrices, the 1st diagonal contains the values 1 and 1 - skipping it.\n",
                "    # Now let's handle the rest of the upper-left triangle.\n",
                "    for skew_diagonal_index in range(2, shorter_dim):\n",
                "        skew_diagonal_length = (skew_diagonal_index + 1)\n",
                "        for offset_in_diagonal in range(1, skew_diagonal_length - 1): # ! Skip the left column & top row\n",
                "            # If we haven't passed the main skew diagonal yet, \n",
                "            # then we have to skip the first and the last operation,\n",
                "            # as those are already pre-populated and form the first column \n",
                "            # and the first row of the Levenshtein matrix respectively.\n",
                "            i = skew_diagonal_index - offset_in_diagonal\n",
                "            j = offset_in_diagonal\n",
                "            if verbose:\n",
                "                print(f\"top left triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            shorter_char = shorter_reversed[len(shorter) - i]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "\n",
                "            # ? For reproducibility let's also try doing the same only using the info in the 3 diagonals\n",
                "            next_distances[offset_in_diagonal] = min(\n",
                "                current_distances[offset_in_diagonal - 1] + 1,\n",
                "                current_distances[offset_in_diagonal] + 1,\n",
                "                previous_distances[offset_in_diagonal - 1] + substitution_cost,\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "            \n",
                "        next_distances[0] = next_distances[skew_diagonal_length-1] = skew_diagonal_index\n",
                "        \n",
                "        # ? Let's validate the contents of the diagonal\n",
                "        skew_diagonal_expected = get_skewed_diagonal(matrix, skew_diagonal_index)\n",
                "        assert len(skew_diagonal_expected) == skew_diagonal_length\n",
                "        assert (skew_diagonal_expected == next_distances[:skew_diagonal_length]).all(), f\"diagonal:{skew_diagonal_index}\\nexpected:{skew_diagonal_expected}\\nproduced:{next_distances[:skew_diagonal_length]}\"\n",
                "        temporary_distances[:] = previous_distances[:]\n",
                "        previous_distances[:] = current_distances[:]\n",
                "        current_distances[:] = next_distances[:]\n",
                "        next_distances[:] = temporary_distances[:]\n",
                "\n",
                "    # Now let's handle the anti-diagonal band of the matrix, between the top and bottom-right triangles.        \n",
                "    for skew_diagonal_index in range(shorter_dim, longer_dim):\n",
                "        skew_diagonal_length = shorter_dim\n",
                "        for offset_in_diagonal in range(skew_diagonal_length - 1): # ! Skip the top row\n",
                "            i = shorter_dim - offset_in_diagonal - 1\n",
                "            j = skew_diagonal_index - shorter_dim + offset_in_diagonal + 1\n",
                "            if verbose:\n",
                "                print(f\"anti-band: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            shorter_char = shorter_reversed[len(shorter) - i]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "    \n",
                "            # ? For reproducibility let's also try doing the same only using the info in the 3 diagonals\n",
                "            next_distances[offset_in_diagonal] = min(\n",
                "                current_distances[offset_in_diagonal] + 1,\n",
                "                current_distances[offset_in_diagonal + 1] + 1,\n",
                "                previous_distances[offset_in_diagonal] + substitution_cost,\n",
                "            )\n",
                "            \n",
                "        next_distances[shorter_dim-1] = skew_diagonal_index\n",
                "        \n",
                "        # ? Let's validate the contents of the diagonal\n",
                "        skew_diagonal_expected = get_skewed_diagonal(matrix, skew_diagonal_index)\n",
                "        assert len(skew_diagonal_expected) == skew_diagonal_length\n",
                "        assert (skew_diagonal_expected == next_distances[:skew_diagonal_length]).all(), f\"diagonal:{skew_diagonal_index}\\nexpected:{skew_diagonal_expected}\\nproduced:{next_distances[:skew_diagonal_length]}\"\n",
                "        temporary_distances[:] = previous_distances[:]\n",
                "        previous_distances[:-1] = current_distances[1:] # ! Note we shift here\n",
                "        current_distances[:] = next_distances[:]\n",
                "        next_distances[:] = temporary_distances[:]\n",
                "    \n",
                "    # Now let's handle the bottom right triangle.\n",
                "    for skew_diagonal_index in range(longer_dim, diagonals_count):\n",
                "        skew_diagonal_length = diagonals_count - skew_diagonal_index\n",
                "        for offset_in_diagonal in range(skew_diagonal_length):\n",
                "            i = shorter_dim - offset_in_diagonal - 1\n",
                "            j = skew_diagonal_index - shorter_dim + offset_in_diagonal + 1\n",
                "            if verbose:\n",
                "                print(f\"bottom right triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            assert (i - 1) >= 0 and (i - 1) < len(shorter), f\"{i = }\"\n",
                "            assert (j - 1) >= 0 and (j - 1) < len(longer), f\"{j = }\"\n",
                "            shorter_char = shorter_reversed[len(shorter) - i]\n",
                "            longer_char = longer[j - 1]\n",
                "            print(f\"{shorter_char=}, {longer_char=}\")\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "\n",
                "            # ? For reproducibility let's also try doing the same only using the info in the 3 diagonals\n",
                "            next_distances[offset_in_diagonal] = min(\n",
                "                current_distances[offset_in_diagonal] + 1,\n",
                "                current_distances[offset_in_diagonal + 1] + 1,\n",
                "                previous_distances[offset_in_diagonal] + substitution_cost,\n",
                "            )\n",
                "\n",
                "        # ? Let's validate the contents of the diagonal\n",
                "        skew_diagonal_expected = get_skewed_diagonal(matrix, skew_diagonal_index)\n",
                "        assert len(skew_diagonal_expected) == skew_diagonal_length\n",
                "        assert (skew_diagonal_expected == next_distances[:skew_diagonal_length]).all(), f\"diagonal:{skew_diagonal_index}\\nexpected:{skew_diagonal_expected}\\nproduced:{next_distances[:skew_diagonal_length]}\"\n",
                "        temporary_distances[:] = previous_distances[:]\n",
                "        previous_distances[:-1] = current_distances[1:] # ! Note we shift here\n",
                "        current_distances[:] = next_distances[:]\n",
                "        next_distances[:] = temporary_distances[:]\n",
                "\n",
                "    # Return the Levenshtein distance\n",
                "    distance_from_matrix = matrix[len(shorter), len(longer)]\n",
                "    distance_from_diagonal = current_distances[0]\n",
                "    assert distance_from_diagonal == distance_from_matrix\n",
                "    if len(s1) > len(s2):\n",
                "        matrix = matrix.T\n",
                "    return distance_from_matrix, matrix"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import random\n",
                "for _ in range(100):\n",
                "    len1 = random.randint(1, 50)\n",
                "    len2 = random.randint(1, 50)\n",
                "    s1 = ''.join(random.choices(\"abc\", k=len1))\n",
                "    s2 = ''.join(random.choices(\"abc\", k=len2))\n",
                "    distance_wf, matrix_wf = wagner_fisher(s1, s2)\n",
                "    distance_sd, matrix_sd = skewed_diagonals_reversed(s1, s2, baseline=matrix_wf, verbose=False)\n",
                "    assert distance_wf == distance_sd, f\"{distance_wf = } != {distance_sd = }\"\n",
                "    assert np.all(matrix_wf == matrix_sd), f\"{matrix_wf = }\\n{matrix_sd = }\""
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "s1 = \"listeners\"\n",
                "s2 = \"silents\"\n",
                "distance_wf, matrix_wf = wagner_fisher(s1, s2)\n",
                "distance_sd, matrix_sd = skewed_diagonals_reversed(s1, s2, baseline=matrix_wf)\n",
                "s1, s2, f\"{distance_wf = }\", f\"{distance_sd = }\", matrix_sd"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "s1 = \"atca\"\n",
                "s2 = \"ctactcaccc\"\n",
                "distance_wf, matrix_wf = wagner_fisher(s1, s2)\n",
                "distance_sd, matrix_sd = skewed_diagonals_reversed(s1, s2, baseline=matrix_wf)\n",
                "s1, s2, f\"{distance_wf = }\", f\"{distance_sd = }\", matrix_sd"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Shift-less with Reverse Order"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from typing import Optional\n",
                "\n",
                "def skewed_diagonals_reversed(\n",
                "    s1: str, s2: str, \n",
                "    verbose: bool = False, \n",
                "    baseline: Optional[np.ndarray] = None) -> Tuple[int, np.ndarray]:\n",
                "    \n",
                "    shorter, longer = (s1, s2) if len(s1) <= len(s2) else (s2, s1)    \n",
                "    baseline = baseline if len(s1) <= len(s2) else baseline.T\n",
                "    shorter_dim = len(shorter) + 1\n",
                "    longer_dim = len(longer) + 1\n",
                "    if verbose:\n",
                "        print(f\"{shorter=}, {longer=}, {shorter_dim=}, {longer_dim=}\")\n",
                "    \n",
                "    # Create a matrix of size (shorter_dim) x (longer_dim)\n",
                "    matrix = np.zeros((shorter_dim, longer_dim), dtype=int)\n",
                "    matrix[:, :] = longer_dim + 1 # or +inf \n",
                "\n",
                "    # Initialize the first column and first row of the matrix\n",
                "    for i in range(shorter_dim):\n",
                "        matrix[i, 0] = i\n",
                "    for j in range(longer_dim):\n",
                "        matrix[0, j] = j\n",
                "\n",
                "    # Let's say we are dealing with 3 and 5 letter words.\n",
                "    # The matrix will have size 4 x 6, parameterized as (shorter_dim x longer_dim).\n",
                "    # It will have:\n",
                "    # - 4 diagonals of increasing length, at positions: 0, 1, 2, 3.\n",
                "    # - 2 diagonals of fixed length, at positions: 4, 5.\n",
                "    # - 3 diagonals of decreasing length, at positions: 6, 7, 8.\n",
                "    diagonals_count = shorter_dim + longer_dim - 1\n",
                "    shorter_reversed = \"\".join(reversed(shorter))\n",
                "\n",
                "    # In reality, we need to keep only 3 diagonals to produce the same score in the end.\n",
                "    previous_distances = np.zeros(shorter_dim, dtype=np.uint)\n",
                "    current_distances = np.zeros(shorter_dim, dtype=np.uint)\n",
                "    next_distances = np.zeros(shorter_dim, dtype=np.uint)\n",
                "    temporary_distances = np.zeros(shorter_dim, dtype=np.uint)\n",
                "    previous_distances[0] = 0\n",
                "    current_distances[0] = current_distances[1] = 1\n",
                "\n",
                "    # Same as with square matrices, the 0th diagonal contains - just one element - zero - skipping it.\n",
                "    # Same as with square matrices, the 1st diagonal contains the values 1 and 1 - skipping it.\n",
                "    # Now let's handle the rest of the upper-left triangle.\n",
                "    for skew_diagonal_index in range(2, shorter_dim):\n",
                "        skew_diagonal_length = (skew_diagonal_index + 1)\n",
                "        for offset_in_diagonal in range(1, skew_diagonal_length - 1): # ! Skip the left column & top row\n",
                "            # If we haven't passed the main skew diagonal yet, \n",
                "            # then we have to skip the first and the last operation,\n",
                "            # as those are already pre-populated and form the first column \n",
                "            # and the first row of the Levenshtein matrix respectively.\n",
                "            i = skew_diagonal_index - offset_in_diagonal\n",
                "            j = offset_in_diagonal\n",
                "            if verbose:\n",
                "                print(f\"top left triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            shorter_char = shorter_reversed[len(shorter) - i]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "\n",
                "            # ? For reproducibility let's also try doing the same only using the info in the 3 diagonals\n",
                "            next_distances[offset_in_diagonal] = min(\n",
                "                current_distances[offset_in_diagonal - 1] + 1,\n",
                "                current_distances[offset_in_diagonal] + 1,\n",
                "                previous_distances[offset_in_diagonal - 1] + substitution_cost,\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "            \n",
                "        next_distances[0] = next_distances[skew_diagonal_length-1] = skew_diagonal_index\n",
                "        \n",
                "        # ? Let's validate the contents of the diagonal\n",
                "        skew_diagonal_expected = get_skewed_diagonal(matrix, skew_diagonal_index)\n",
                "        assert len(skew_diagonal_expected) == skew_diagonal_length\n",
                "        assert (skew_diagonal_expected == next_distances[:skew_diagonal_length]).all(), f\"diagonal:{skew_diagonal_index}\\nexpected:{skew_diagonal_expected}\\nproduced:{next_distances[:skew_diagonal_length]}\"\n",
                "        temporary_distances[:] = previous_distances[:]\n",
                "        previous_distances[:] = current_distances[:]\n",
                "        current_distances[:] = next_distances[:]\n",
                "        next_distances[:] = temporary_distances[:]\n",
                "\n",
                "    # Now let's handle the anti-diagonal band of the matrix, between the top and bottom-right triangles.        \n",
                "    for skew_diagonal_index in range(shorter_dim, longer_dim):\n",
                "        skew_diagonal_length = shorter_dim\n",
                "        for offset_in_diagonal in range(skew_diagonal_length - 1): # ! Skip the top row\n",
                "            i = shorter_dim - offset_in_diagonal - 1\n",
                "            j = skew_diagonal_index - shorter_dim + offset_in_diagonal + 1\n",
                "            if verbose:\n",
                "                print(f\"anti-band: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            shorter_char = shorter_reversed[len(shorter) - i]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "    \n",
                "            # ? For reproducibility let's also try doing the same only using the info in the 3 diagonals\n",
                "            next_distances[offset_in_diagonal] = min(\n",
                "                current_distances[offset_in_diagonal] + 1,\n",
                "                current_distances[offset_in_diagonal + 1] + 1,\n",
                "                previous_distances[offset_in_diagonal] + substitution_cost,\n",
                "            )\n",
                "            \n",
                "        next_distances[shorter_dim-1] = skew_diagonal_index\n",
                "        \n",
                "        # ? Let's validate the contents of the diagonal\n",
                "        skew_diagonal_expected = get_skewed_diagonal(matrix, skew_diagonal_index)\n",
                "        assert len(skew_diagonal_expected) == skew_diagonal_length\n",
                "        assert (skew_diagonal_expected == next_distances[:skew_diagonal_length]).all(), f\"diagonal:{skew_diagonal_index}\\nexpected:{skew_diagonal_expected}\\nproduced:{next_distances[:skew_diagonal_length]}\"\n",
                "        temporary_distances[:] = previous_distances[:]\n",
                "        previous_distances[:-1] = current_distances[1:] # ! Note we shift here\n",
                "        current_distances[:] = next_distances[:]\n",
                "        next_distances[:] = temporary_distances[:]\n",
                "    \n",
                "    # Now let's handle the bottom right triangle.\n",
                "    for skew_diagonal_index in range(longer_dim, diagonals_count):\n",
                "        skew_diagonal_length = diagonals_count - skew_diagonal_index\n",
                "        for offset_in_diagonal in range(skew_diagonal_length):\n",
                "            i = shorter_dim - offset_in_diagonal - 1\n",
                "            j = skew_diagonal_index - shorter_dim + offset_in_diagonal + 1\n",
                "            if verbose:\n",
                "                print(f\"bottom right triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\")\n",
                "            assert (i - 1) >= 0 and (i - 1) < len(shorter), f\"{i = }\"\n",
                "            assert (j - 1) >= 0 and (j - 1) < len(longer), f\"{j = }\"\n",
                "            shorter_char = shorter_reversed[len(shorter) - i]\n",
                "            longer_char = longer[j - 1]\n",
                "            print(f\"{shorter_char=}, {longer_char=}\")\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,                      # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,                      # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "            \n",
                "            if baseline is not None:\n",
                "                assert matrix[i, j] == baseline[i, j], f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "\n",
                "            # ? For reproducibility let's also try doing the same only using the info in the 3 diagonals\n",
                "            next_distances[offset_in_diagonal] = min(\n",
                "                current_distances[offset_in_diagonal] + 1,\n",
                "                current_distances[offset_in_diagonal + 1] + 1,\n",
                "                previous_distances[offset_in_diagonal] + substitution_cost,\n",
                "            )\n",
                "\n",
                "        # ? Let's validate the contents of the diagonal\n",
                "        skew_diagonal_expected = get_skewed_diagonal(matrix, skew_diagonal_index)\n",
                "        assert len(skew_diagonal_expected) == skew_diagonal_length\n",
                "        assert (skew_diagonal_expected == next_distances[:skew_diagonal_length]).all(), f\"diagonal:{skew_diagonal_index}\\nexpected:{skew_diagonal_expected}\\nproduced:{next_distances[:skew_diagonal_length]}\"\n",
                "        temporary_distances[:] = previous_distances[:]\n",
                "        previous_distances[:-1] = current_distances[1:] # ! Note we shift here\n",
                "        current_distances[:] = next_distances[:]\n",
                "        next_distances[:] = temporary_distances[:]\n",
                "\n",
                "    # Return the Levenshtein distance\n",
                "    distance_from_matrix = matrix[len(shorter), len(longer)]\n",
                "    distance_from_diagonal = current_distances[0]\n",
                "    assert distance_from_diagonal == distance_from_matrix\n",
                "    if len(s1) > len(s2):\n",
                "        matrix = matrix.T\n",
                "    return distance_from_matrix, matrix"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Bounding the Error\n",
                "\n",
                "It's easy to spot that the algorithm can be further optimized if we are dealing with \"bounded\" edit distances, where the maximum allowed number of edits is known in advance.\n",
                "In such cases, we only need to evaluate a band around the main diagonal, and can skip the rest of the matrix.\n",
                "For the bound $k$, we only need to evaluate $2k+1$ diagonals."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from typing import Optional\n",
                "\n",
                "\n",
                "def bounded_skewed_diagonals(\n",
                "    s1: str,\n",
                "    s2: str,\n",
                "    verbose: bool = False,\n",
                "    bound: Optional[int] = None,\n",
                "    baseline: Optional[np.ndarray] = None,\n",
                ") -> Tuple[int, np.ndarray]:\n",
                "\n",
                "    shorter, longer = (s1, s2) if len(s1) <= len(s2) else (s2, s1)\n",
                "    baseline = baseline if len(s1) <= len(s2) else baseline.T\n",
                "    shorter_dim = len(shorter) + 1\n",
                "    longer_dim = len(longer) + 1\n",
                "    if verbose:\n",
                "        print(f\"{shorter=}, {longer=}, {shorter_dim=}, {longer_dim=}\")\n",
                "\n",
                "    # Create a matrix of size (shorter_dim) x (longer_dim)\n",
                "    matrix = np.zeros((shorter_dim, longer_dim), dtype=int)\n",
                "    matrix[:, :] = np.iinfo(matrix.dtype).max\n",
                "\n",
                "    # Initialize the first column and first row of the matrix\n",
                "    for i in range(shorter_dim):\n",
                "        matrix[i, 0] = i\n",
                "    for j in range(longer_dim):\n",
                "        matrix[0, j] = j\n",
                "\n",
                "    # Let's say we are dealing with 3 and 5 letter words.\n",
                "    # The matrix will have size 4 x 6, parameterized as (shorter_dim x longer_dim).\n",
                "    # It will have:\n",
                "    # - 4 diagonals of increasing length, at positions: 0, 1, 2, 3.\n",
                "    # - 2 diagonals of fixed length, at positions: 4, 5.\n",
                "    # - 3 diagonals of decreasing length, at positions: 6, 7, 8.\n",
                "    diagonals_count = shorter_dim + longer_dim - 1\n",
                "\n",
                "    # Same as with square matrices, the 0th diagonal contains - just one element - zero - skipping it.\n",
                "    # Same as with square matrices, the 1st diagonal contains the values 1 and 1 - skipping it.\n",
                "    # In unbounded case, we the upper-left triangle will have `shorter_dim` rows and columns.\n",
                "    # In bounded case, we will have `min(bound, shorter_dim)` rows and columns.\n",
                "    upper_triangle_dim = min(bound, shorter_dim) if bound is not None else shorter_dim\n",
                "    for skew_diagonal_index in range(2, upper_triangle_dim):\n",
                "        skew_diagonal_length = skew_diagonal_index + 1\n",
                "        for offset_within_diagonal in range(\n",
                "            1, skew_diagonal_length - 1\n",
                "        ):  #! Skip the first column & row\n",
                "            # If we haven't passed the main skew diagonal yet,\n",
                "            # then we have to skip the first and the last operation,\n",
                "            # as those are already pre-populated and form the first column\n",
                "            # and the first row of the Levenshtein matrix respectively.\n",
                "            i = skew_diagonal_index - offset_within_diagonal\n",
                "            j = offset_within_diagonal\n",
                "            if verbose:\n",
                "                print(\n",
                "                    f\"top left triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\"\n",
                "                )\n",
                "            shorter_char = shorter[i - 1]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,  # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,  # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "            \n",
                "            # Validation checks:\n",
                "            if baseline is not None:\n",
                "                assert (\n",
                "                    matrix[i, j] == baseline[i, j]\n",
                "                ), f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "\n",
                "    # Now let's handle the anti-diagonal band of the matrix, between the top and bottom-right triangles.\n",
                "    # In the unbounded case, we will enumerate diagonal indices from `shorter_dim` to `longer_dim`.\n",
                "    # In the bounded case, we go through the same \n",
                "    for skew_diagonal_index in range(shorter_dim, longer_dim):\n",
                "        skew_diagonal_length = shorter_dim\n",
                "        for offset_within_diagonal in range(\n",
                "            skew_diagonal_length - 1\n",
                "        ):  #! Skip the first row\n",
                "            i = shorter_dim - offset_within_diagonal - 1\n",
                "            j = skew_diagonal_index - shorter_dim + offset_within_diagonal + 1\n",
                "            if verbose:\n",
                "                print(\n",
                "                    f\"anti-band: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\"\n",
                "                )\n",
                "            shorter_char = shorter[i - 1]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,  # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,  # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "\n",
                "            if baseline is not None:\n",
                "                assert (\n",
                "                    matrix[i, j] == baseline[i, j]\n",
                "                ), f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "\n",
                "    # Now let's handle the bottom right triangle.\n",
                "    for skew_diagonal_index in range(longer_dim, diagonals_count):\n",
                "        skew_diagonal_length = diagonals_count - skew_diagonal_index\n",
                "        for offset_within_diagonal in range(skew_diagonal_length):\n",
                "            i = shorter_dim - offset_within_diagonal - 1\n",
                "            j = skew_diagonal_index - shorter_dim + offset_within_diagonal + 1\n",
                "            if verbose:\n",
                "                print(\n",
                "                    f\"bottom right triangle: {skew_diagonal_index=}, {skew_diagonal_length=}, {i=}, {j=}\"\n",
                "                )\n",
                "            assert (i - 1) >= 0 and (i - 1) < len(shorter), f\"{i = }\"\n",
                "            assert (j - 1) >= 0 and (j - 1) < len(longer), f\"{j = }\"\n",
                "            shorter_char = shorter[i - 1]\n",
                "            longer_char = longer[j - 1]\n",
                "            substitution_cost = shorter_char != longer_char\n",
                "            matrix[i, j] = min(\n",
                "                matrix[i - 1, j] + 1,  # ? Deletion cost\n",
                "                matrix[i, j - 1] + 1,  # ? Insertion cost\n",
                "                matrix[i - 1, j - 1] + substitution_cost,  # ? Substitution cost\n",
                "            )\n",
                "\n",
                "            if baseline is not None:\n",
                "                assert (\n",
                "                    matrix[i, j] == baseline[i, j]\n",
                "                ), f\"{matrix[i, j]} != {baseline[i, j]} at {i=}, {j=}\"\n",
                "\n",
                "    # Return the Levenshtein distance\n",
                "    distance = matrix[len(shorter), len(longer)]\n",
                "    if len(s1) > len(s2):\n",
                "        matrix = matrix.T\n",
                "    return distance, matrix"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Putting Everything Together"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def vectorized_skewed_diagonals(\n",
                "    s1: str, s2: str, \n",
                "    verbose: bool = False, \n",
                "    baseline: Optional[np.ndarray] = None) -> Tuple[int, np.ndarray]:\n",
                "    \n",
                "    shorter, longer = (s1, s2) if len(s1) <= len(s2) else (s2, s1)    \n",
                "    baseline = baseline if len(s1) <= len(s2) else baseline.T\n",
                "    shorter_dim = len(shorter) + 1\n",
                "    longer_dim = len(longer) + 1\n",
                "    if verbose:\n",
                "        print(f\"{shorter=}, {longer=}, {shorter_dim=}, {longer_dim=}\")\n",
                "    \n",
                "    # Create a matrix of size (shorter_dim) x (longer_dim)\n",
                "    matrix = np.zeros((shorter_dim, longer_dim), dtype=int)\n",
                "    matrix[:, :] = longer_dim + 1 # or +inf \n",
                "\n",
                "    # Initialize the first column and first row of the matrix\n",
                "    for i in range(shorter_dim):\n",
                "        matrix[i, 0] = i\n",
                "    for j in range(longer_dim):\n",
                "        matrix[0, j] = j\n",
                "\n",
                "    # Let's say we are dealing with 3 and 5 letter words.\n",
                "    # The matrix will have size 4 x 6, parameterized as (shorter_dim x longer_dim).\n",
                "    # It will have:\n",
                "    # - 4 diagonals of increasing length, at positions: 0, 1, 2, 3.\n",
                "    # - 2 diagonals of fixed length, at positions: 4, 5.\n",
                "    # - 3 diagonals of decreasing length, at positions: 6, 7, 8.\n",
                "    diagonals_count = shorter_dim + longer_dim - 1\n",
                "\n",
                "    # Same as with square matrices, the 0th diagonal contains - just one element - zero - skipping it.\n",
                "    # Same as with square matrices, the 1st diagonal contains the values 1 and 1 - skipping it.\n",
                "    # Now let's handle the rest of the upper-left triangle.\n",
                "    next_diagonal_index = 2\n",
                "    while next_diagonal_index < shorter_dim:\n",
                "        next_skew_diagonal_length = next_diagonal_index + 1\n",
                "\n",
                "        old_substitution_costs = previous[:next_skew_diagonal_length - 2]\n",
                "        added_substitution_costs = [shorter[next_diagonal_index - offset_within_diagonal - 2] != longer[offset_within_diagonal] for offset_within_diagonal in range(next_skew_diagonal_length - 2)]\n",
                "        substitution_costs = old_substitution_costs + added_substitution_costs\n",
                "\n",
                "        following[1:next_skew_diagonal_length - 1] = np.minimum(current[1:next_skew_diagonal_length - 1] + 1, current[:next_skew_diagonal_length - 2] + 1) # Insertions or deletions\n",
                "        following[1:next_skew_diagonal_length - 1] = np.minimum(following[1:next_skew_diagonal_length - 1], substitution_costs) # Substitutions\n",
                "        following[0] = next_diagonal_index\n",
                "        following[next_skew_diagonal_length - 1] = next_diagonal_index\n",
                "        assert np.all(following[:next_skew_diagonal_length] == get_skewed_diagonal(baseline, next_diagonal_index))\n",
                "        \n",
                "        previous[:] = current[:]\n",
                "        current[:] = following[:]\n",
                "        next_diagonal_index += 1\n",
                "                        \n",
                "    # Now let's handle the anti-diagonal band of the matrix, between the top and bottom-right triangles.        \n",
                "    while next_diagonal_index < longer_dim:\n",
                "        next_skew_diagonal_length = shorter_dim\n",
                "    \n",
                "        ..."
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "StringZilla",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.11.11"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}
